from channels.db import database_sync_to_async
from channels.generic.websocket import AsyncJsonWebsocketConsumer
from django.db import transaction
from rest_framework_jwt.serializers import VerifyJSONWebTokenSerializer

from apps.chat.serializers import ChatMessageModelSerializer
from apps.users.models import ChatClients, ChatMessage, User


@database_sync_to_async
def verify_token(token):
    return VerifyJSONWebTokenSerializer().validate(token)


class TokenAuthMiddlewareStack:

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        headers = dict(scope['headers'])
        token = {}
        if b'sec-websocket-protocol' in headers:
            token = {'token': headers[b'sec-websocket-protocol']}
        try:
            token_obj = await verify_token(token)
        except Exception as e:
            print(e)
            # 解析失败，则不做任何处理，直接交给视图函数，视图函数会尝试从scope中取出关键信息，失败则会断开websocket连接
            return await self.app(scope, receive, send)
        # 将客户端的唯一身份标识（关键信息）加入scope
        scope['user'] = token_obj['user']
        scope['token'] = token_obj['token'].decode()
        return await self.app(scope, receive, send)


class AsyncChatConsumers(AsyncJsonWebsocketConsumer):
    def __init__(self):
        super().__init__()
        self.online_group = 'online_users'

    async def websocket_connect(self, message):
        if not self.scope.get('token', None):
            return
        await self.base_send({"type": "websocket.accept", "subprotocol": self.scope['token']})
        await self.init_client()
        # Join room group
        await self.channel_layer.group_add(
            self.online_group,
            self.channel_name
        )

    async def connect(self):
        await self.accept()

    async def disconnect(self, close_code):
        # Leave room group
        await self.channel_layer.group_discard(
            self.online_group,
            self.channel_name
        )

    # Receive message from WebSocket
    async def receive_json(self, content, **kwargs):
        msg_type = content.get('type')
        receiver_id = content.get('receiver_id')
        message = content.get('message')
        receiver = await self.get_receiver(receiver_id)
        if not receiver:
            return
        if msg_type == 'chat.message':
            await self.send_chat_message(receiver, message)
        # # Send message to room group
        # await self.channel_layer.group_send(
        #     self.online_group,
        #     {
        #         'type': 'chat_message',  # type 用于指定该消息的类型，根据消息类型调用不同的函数去处理消息
        #         'message': message  # message 内为消息主体
        #     }
        # )

    # Receive message from room group
    async def chat_message(self, event):
        await self.send_json({'message': event['message']})

    async def system_message(self, event):
        await self.send_json({'message': event['message']})

    @database_sync_to_async
    def init_client(self):
        try:
            ChatClients.objects.create(user=self.scope['user'], channel_name=self.channel_name)
        except Exception as e:
            ChatClients.objects.filter(user=self.scope['user']).update(channel_name=self.channel_name)

    @database_sync_to_async
    def get_receiver(self, receiver_id):
        try:
            return ChatClients.objects.get(user_id=receiver_id)
        except Exception as e:
            print(e)
            return 0

    @database_sync_to_async
    def create_chat_message(self, receiver, message):
        try:
            message = ChatMessage.objects.create(sender=self.scope['user'], receiver=receiver.user, message=message)
            return message
        except Exception as e:
            print(e)
            return 0

    async def send_chat_message(self, receiver, message):
        message_obj = await self.create_chat_message(receiver, message)
        serializer = ChatMessageModelSerializer(message_obj)
        if not message_obj:
            return
        await self.channel_layer.send(receiver.channel_name, {'type': 'chat.message', 'message': serializer.data})
        await self.channel_layer.send(self.channel_name, {'type': 'chat.message', 'message': serializer.data})
